from numpy.random.mtrand import rand, random, uniform
from base import ItemInterface, StudentInterface, ModelCompute, BaseQ
import numpy as np


class StudentParam(StudentInterface):

    def __init__(self, student_param_mean, student_param_cov, n, beta_theta, beta_tau):
        self.n = n
        self.student_param_mean = student_param_mean
        self.student_param_cov = student_param_cov
        self.beta_theta = beta_theta
        self.beta_tau = beta_tau
        self.theta = 0

        self.tau = 0

    def get(self):
        self.student_param = np.random.multivariate_normal(self.student_param_mean,
                                                           self.student_param_cov,
                                                           size=self.n)
        self.theta = self.student_param[:, 0]
        self.tau = self.student_param[:, 1]

        return self.theta, self.tau


class IrtBaseModel(ItemInterface):
    def __init__(self, i, k, a_param, b_xi_param, c_param, alpha_param):
        super().__init__(i)
        self.k = k
        self.a_param = a_param
        self.b_xi_param = b_xi_param
        self.c_param = c_param
        self.alpha_param = alpha_param
        self.a = None
        self.b = None
        self.c = None
        self.xi = None
        self.alpha = None

    def get(self, *args, **kwargs):
        self.a = np.random.uniform(self.a_param[0],
                                   self.a_param[1], size=self.i)
        b_xi = np.random.multivariate_normal(mean=self.b_xi_param[0],
                                             cov=self.b_xi_param[1],
                                             size=self.i)
        self.b = b_xi[:, 0]
        self.xi = b_xi[:, 1]
        self.c = np.random.uniform(self.c_param[0],
                                   self.c_param[1], size=self.i)
        self.alpha = np.random.uniform(self.alpha_param[0],
                                       self.alpha_param[1], size=self.i)


class JointModelCompute(ModelCompute):
    def __init__(self, i, k, n, item_class, student_class):
        super().__init__(i, k, n)
        self.item_class = item_class
        self.student_class = student_class
        self.ra = None
        self.rt = None
        self.ra_possibility = None
        self.rt_mean = None
        self.rt_sigma = None
        self.Q = BaseQ(self.i, self.k, R_cnt=2).get()

    def get(self):
        if self.ra_possibility == None:
            self.get_possibility()
        self.ra = np.random.binomial(1, self.ra_possibility)

    def get_possibility(self):
        self.ra_possibility = self.ra_model()
        self.rt_model()

    def ra_model(self):
        _2pl = self.student_class.theta * self.item_class.a - self.item_class.b
        _3pl = self.item_class.c + (1 - self.item_class.c) * _2pl
        return np.exp(_3pl) / (1 + np.exp(_3pl))

    def rt_model(self):
        rt_mean = self.item_class.xi - self.student_class.tau.dot(self.Q.T)
        rt_scale = np.stack([1 / self.item_class.alpha for i in range(self.student_class.n)])
        self.rt_mean = rt_mean
        self.rt_sigma = rt_scale ** 2
        self.rt = np.exp(np.random.normal(loc=rt_mean, scale=rt_scale))


class AberrantResponse(object):
    def __init__(self, n, i, student_class, item_class, compute_class):
        self.n = n
        self.i = i
        self.student_class = student_class
        self.item_class = item_class
        self.compute_class = compute_class
        self.rg_flag = None
        self.preknowledge_flag = np.zeros((self.n, self.i))
        self.d_c = None
        self.d_g = None
        self.ra = None
        self.rt = None

    def get(self, rg_param, cheat_param):
        cheat_rap = np.random.uniform(cheat_param[0], cheat_param[1], size=self.i)
        rg_rap = np.random.beta(rg_param[0], rg_param[1], size=self.i)
        cheat_resp = np.random.binomial(1, cheat_rap, size=(self.n, self.i))
        rg_resp = np.random.binomial(1, rg_rap, size=(self.n, self.i))
        self.d_c = cheat_rap
        self.d_g = rg_rap

        # 生成结果同时为作弊和快速猜测
        rg_and_cheating = (self.rg_flag == 1) * (self.preknowledge_flag == 1)
        # 按照0.5的概率随机分到作弊或者快速猜测
        re_sample = np.random.binomial(1, rg_and_cheating * 0.5)
        self.rg_flag = self.rg_flag - rg_and_cheating + re_sample * rg_and_cheating
        self.preknowledge_flag = self.preknowledge_flag - rg_and_cheating + (1 - re_sample) * rg_and_cheating
        # 生成最终作答
        self.ra = (1 - self.rg_flag - self.preknowledge_flag) * self.compute_class.ra \
                  + self.rg_flag * rg_resp + \
                  self.preknowledge_flag * cheat_resp
        self.rt = (1 - self.rg_flag - self.preknowledge_flag) * self.compute_class.rt + \
                  np.exp(np.random.normal(-2, 0.25, size=(self.n, self.i))) * (self.rg_flag + self.preknowledge_flag)

    def flag_rg(self, rg_portion, rg_severity):
        stuent_percent = np.percentile(self.student_class.tau, [1 / 3, 2 / 3])
        percent_mat = (self.student_class.tau > stuent_percent[0]) \
                      + (self.student_class.tau > stuent_percent[1])
        x = rg_portion * 3
        group_rg_possibility = x * np.array([0.6, 0.3, 0.1])
        stuent_level_rg_p = (percent_mat == 0 * group_rg_possibility[0]) + \
                            (percent_mat == 1 * group_rg_possibility[1]) + \
                            (percent_mat == 2 * group_rg_possibility[2])
        print(group_rg_possibility)
        student_item_rg_p = np.outer(stuent_level_rg_p, np.repeat(rg_severity, self.i))
        self.rg_flag = np.random.binomial(1, student_item_rg_p)

    def flag_preknowledge(self, compromise_item_index, severity):
        flag = np.zeros((self.n, self.i))
        flag[:, compromise_item_index] = np.random.binomial(1, severity, size=(self.n, len(compromise_item_index)))
        self.preknowledge_flag = flag


class CdmStudentParam(StudentInterface):
    def __init__(self, n, k, student_param_mean, student_param_cov, beta_theta, delta_att):
        super().__init__(n, k)
        self.student_param_mean = student_param_mean
        self.student_param_cov = student_param_cov
        self.beta_theta = beta_theta
        self.delta_att = delta_att
        self.theta = 0
        self.specific_att = None
        self.tau = 0

    def get(self):
        self.student_param = np.random.multivariate_normal(self.student_param_mean,
                                                           self.student_param_cov,
                                                           size=self.n)
        self.theta = self.student_param[:, 0]
        self.tau = self.student_param[:, 1]
        self.specific_att = self.get_specific_att()
        return self.theta, self.tau, self.specific_att

    def get_specific_att(self):
        _ = np.outer(self.theta, self.beta_theta) + self.delta_att
        att_possibility = np.exp(_) / (1 + np.exp(_))
        return np.random.binomial(1, att_possibility)


class CdmBaseModel(ItemInterface):

    def __init__(self, i, k, mean_param, cov_param, alpha_param):
        super().__init__(i)
        self.k = k
        self.item_mean = mean_param
        self.item_cov = cov_param
        self.alpha_param = alpha_param
        self.item_param = None
        self.beta_i = None
        self.delta_i = None
        self.xi = None

    def get(self):
        self.item_param = np.random.multivariate_normal(self.item_mean,
                                                        self.item_cov,
                                                        size=self.i)
        self.beta_i = self.item_param[:, 0]
        self.delta_i = self.item_param[:, 1]
        self.xi = self.item_param[:, 2]

        self.alpha = np.random.uniform(self.alpha_param[0],
                                       self.alpha_param[1], size=self.i)

        self.trans_to_s_g()

    def trans_to_s_g(self):
        self.g = np.exp(self.beta_i) / (1 + np.exp(self.beta_i))
        self.s = 1 - np.exp(self.delta_i + self.beta_i) / (1 + np.exp(self.delta_i + self.beta_i))


class CdmJointModelCompute(ModelCompute):
    def __init__(self, i, k, n, item_class, student_class):
        super().__init__(i, k, n)
        self.item_class = item_class
        self.student_class = student_class
        self.ra = None
        self.rt = None
        self.ra_possibility = None
        self.rt_mean = None
        self.rt_sigma = None
        self.Q = BaseQ(self.i, self.k, R_cnt=2).get()

    def get(self):
        if self.ra_possibility == None:
            self.get_possibility()
        self.ra = np.random.binomial(1, self.ra_possibility)

    def get_possibility(self):
        self.ra_model()
        self.rt_model()

    def ra_model(self):
        eta = (self.student_class.specific_att.reshape(self.n, 1, self.k) >= \
               self.Q.reshape(1, self.i, self.k)).prod(axis=-1)
        self.ra_possibility = self.item_class.g + (1 - self.item_class.s - self.item_class.g) * eta
        self.ra = np.random.binomial(1, self.ra_possibility)

    def rt_model(self):
        rt_mean = self.item_class.xi.reshape(1, -1) - self.student_class.tau.reshape(-1, 1)
        rt_scale = np.stack([1 / self.item_class.alpha for i in range(self.student_class.n)])
        self.rt_mean = rt_mean
        self.rt_sigma = rt_scale ** 2
        self.rt = np.exp(np.random.normal(loc=rt_mean, scale=rt_scale))


class Abberant():
    def __init__(self, n, i, student_class, item_class, compute_class, d_g, d_c, aberrant_rt_mean, aberrant_rt_var):
        self.n = n
        self.i = i
        self.student_class = student_class
        self.item_class = item_class
        self.compute_class = compute_class
        self.root_ra = compute_class.ra
        self.root_rt = compute_class.rt
        self.flag = np.zeros((self.n, self.i))
        self.student_flag = np.zeros(n)
        self.aberrant_rt_mean = aberrant_rt_mean
        self.aberrant_rt_sigma = aberrant_rt_var ** 0.5
        self.d_c = d_c
        self.d_g = d_g
        self.ra = compute_class.ra
        self.rt = compute_class.rt

    def time_limit(self, time):
        _ = np.zeros_like(self.rt)
        for n in range(len(self.rt)):
            idx = 0
            while (np.sum(self.rt[n]) > time) and (idx > -self.i):
                idx -= 1
                if self.flag[n, idx] == 0:
                    # 标记RG
                    self.flag[n, idx] = 1
                    # 替换RT
                    self.rt[n, idx] = np.exp(np.random.normal(self.aberrant_rt_mean, self.aberrant_rt_sigma))
                    _[n, idx] = 1

        # 替换RA
        self.ra = _ * np.random.binomial(1, self.d_g, size=self.ra.shape) \
                  + (1 - _) * self.root_ra

    def random_rg(self, prob, severity):
        random_flag = np.zeros_like(self.rt)
        p_index = np.random.binomial(1, prob, self.n)
        cnt = -1
        for _ in p_index:
            cnt += 1
            if _ == 1:
                # self.student_flag[cnt] = 1
                person_flag = np.random.binomial(1, severity, size=self.i)
                for i in range(len(person_flag)):
                    if self.flag[cnt, i] == 0:

                        # flag RG
                        self.flag[cnt, i] = person_flag[i]
                        random_flag[cnt, i] = person_flag[i]
                        # if random_flag[cnt, i] == 1:
                        #     self.rt[cnt, i] = np.exp(np.random.normal(self.aberrant_rt_mean, self.aberrant_rt_sigma))
                        #     self.ra[cnt, i] = np.random.binomial(1, self.d_g)
        self.random_flag = random_flag
        # 替换RT
        self.rt=random_flag*np.exp(np.random.normal(self.aberrant_rt_mean,self.aberrant_rt_sigma,size=(self.rt.shape)))\
                +(1-random_flag)*self.rt

        # 替换RA
        self.ra = random_flag*np.random.binomial(1,self.d_g,size=self.ra.shape)\
                +(1-random_flag)*self.ra

    def random_preknowledge(self, compromise_item_index, severity):
        all_flag = np.random.binomial(1, severity, size=(self.rt.shape))
        for i in range(all_flag.shape[1]):
            if i not in compromise_item_index:
                all_flag[:, i] = 0
        
        for n in range(len(self.student_flag)):
            if self.student_flag[n] != 0:
                all_flag[n,:] = 0


        # flag preknowledge
        self.flag = all_flag * all_flag * 2 + (1 - all_flag) * self.flag

        # flag RT
        self.rt = all_flag * np.exp(
            np.random.normal(self.aberrant_rt_mean, self.aberrant_rt_sigma, size=(self.rt.shape))) + \
                  (1 - all_flag) * self.rt
        # flag RA
        self.ra = all_flag * np.random.binomial(1, self.d_c) + (1 - all_flag) * self.ra


def update():
    return "8"